from ast import main
from ..build_embeddings import EmbeddingBuilder, get_detailed_instruct
import os
import csv

from zkl_pyutils_fsspec import FsLike, resolve_fs
from zkl_pyutils_promptui import confirm_overwrite_path

from ..io.omics_dataset_writer import OmicsDatasetWriter

class EmbeddingWriter:
    def __init__(self, path, model_name="Qwen/Qwen3-Embedding-0.6B", task_description=None, tensor_parallel_size=2, check_existing_keys=True):
        # 确认目标目录是否清空
        confirm_overwrite_path(path)
        
        self.builder = EmbeddingBuilder(model_name, tensor_parallel_size=tensor_parallel_size)
        self.path = path
        self.dataset_path = os.path.join(self.path, "dataset")

        # 若path不存在，则创建
        if not os.path.exists(self.path):
            os.makedirs(self.path)

        self.task_description = task_description
        self.check_existing_keys = check_existing_keys

        self.omics_dataset_init()
        
        self.csv_init()

    def csv_init(self):
        """初始化CSV文件写入"""
        self.csv_file_path = os.path.join(self.path, "str_list.csv")
        self.csv_file = open(self.csv_file_path, 'w', newline='', encoding='utf-8')
        self.csv_writer = csv.writer(self.csv_file)
        self.csv_writer.writerow(['key'])

        
    # 初始化OmicsDataset，返回可写入的对象
    def omics_dataset_init(self):
        # 确认目标目录是否清空
        # confirm_overwrite_path(self.path)

        dst = resolve_fs(self.dataset_path)


        self.dataset_writer = OmicsDatasetWriter(dst)
        self.str_embeddings_writer = self.dataset_writer.str_embeddings_writer()
        self.str_embeddings_writer.write_columns_name(['embeddings'])
    

    def embedding_and_write_batch(self, texts):
        # 对texts进行embedding
        embeddings = self.builder.embed(texts)

        # 循环写入
        for text, embedding in zip(texts, embeddings):
            self.csv_writer.writerow([text])
            self.str_embeddings_writer.write([embedding])
    
    def embedding_and_write(self, text):
        self.embedding_and_write_batch([text])
    
    def save_list(self):
        pass

    def close(self):
        self.builder.close()
        self.str_embeddings_writer.close()
        self.dataset_writer.close()
        if hasattr(self, 'csv_file') and self.csv_file and not self.csv_file.closed:
            self.csv_file.close()

    def __del__(self):
        self.close()
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()
    




      
    
        






        